#include "client/websocket_client.h"
#include "common/utils.h"
#include "common/websocket_frame.h"
#include <iostream>
#include <cstring>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>

namespace websocket {

WebSocketClient::WebSocketClient(const std::string& host, int port)
    : host_(host), port_(port), socket_(-1), connected_(false) {
}

WebSocketClient::~WebSocketClient() {
    disconnect();
}

bool WebSocketClient::connect() {
    // 创建socket
    socket_ = socket(AF_INET, SOCK_STREAM, 0);
    if (socket_ < 0) {
        std::cerr << "创建socket失败" << std::endl;
        return false;
    }
    
    // 解析主机名
    sockaddr_in server_addr{};
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port_);
    
    if (inet_pton(AF_INET, host_.c_str(), &server_addr.sin_addr) <= 0) {
        std::cerr << "无效的服务器地址: " << host_ << std::endl;
        return false;
    }
    
    // 连接服务器
    if (::connect(socket_, reinterpret_cast<sockaddr*>(&server_addr), sizeof(server_addr)) < 0) {
        std::cerr << "连接服务器失败: " << host_ << ":" << port_ << std::endl;
        return false;
    }
    
    // 发送WebSocket握手请求
    std::string handshake = createHandshakeRequest();
    if (!sendData(handshake.c_str(), handshake.length())) {
        std::cerr << "发送握手请求失败" << std::endl;
        return false;
    }
    
    // 接收握手响应
    char buffer[4096];
    ssize_t bytes_read = recv(socket_, buffer, sizeof(buffer) - 1, 0);
    if (bytes_read <= 0) {
        std::cerr << "接收握手响应失败" << std::endl;
        return false;
    }
    
    buffer[bytes_read] = '\0';
    std::string response(buffer);
    
    // 验证握手响应
    if (response.find("HTTP/1.1 101") == std::string::npos ||
        response.find("Upgrade: websocket") == std::string::npos) {
        std::cerr << "WebSocket握手失败" << std::endl;
        return false;
    }
    
    connected_ = true;
    read_thread_ = std::thread(&WebSocketClient::readLoop, this);
    
    std::cout << "成功连接到WebSocket服务器: " << host_ << ":" << port_ << std::endl;
    return true;
}

void WebSocketClient::disconnect() {
    if (connected_) {
        connected_ = false;
        
        if (socket_ >= 0) {
            // 发送关闭帧
            WebSocketFrame close_frame = WebSocketFrame::createCloseFrame();
            auto data = close_frame.serialize();
            sendData(reinterpret_cast<const char*>(data.data()), data.size());
            
            close(socket_);
            socket_ = -1;
        }
        
        if (read_thread_.joinable()) {
            read_thread_.join();
        }
    }
}

void WebSocketClient::sendText(const std::string& message) {
    if (!connected_) return;
    
    WebSocketFrame frame = WebSocketFrame::createTextFrame(message);
    auto data = frame.serialize();
    sendData(reinterpret_cast<const char*>(data.data()), data.size());
}

void WebSocketClient::setMessageHandler(MessageHandler handler) {
    message_handler_ = handler;
}

std::string WebSocketClient::createHandshakeRequest() {
    std::string key = "dGhlIHNhbXBsZSBub25jZQ=="; // 示例key
    
    std::string request = 
        "GET / HTTP/1.1\r\n"
        "Host: " + host_ + ":" + std::to_string(port_) + "\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        "Sec-WebSocket-Key: " + key + "\r\n"
        "Sec-WebSocket-Version: 13\r\n"
        "User-Agent: WebSocketClient/1.0\r\n"
        "\r\n";
    
    return request;
}

void WebSocketClient::readLoop() {
    while (connected_) {
        // 读取帧头
        char header[2];
        if (!readData(header, 2)) {
            break;
        }
        
        // 解析帧头
        uint8_t second_byte = header[1];
        bool masked = (second_byte & 0x80) != 0;
        uint64_t payload_len = second_byte & 0x7F;
        
        // 处理扩展长度
        if (payload_len == 126) {
            char ext_len[2];
            if (!readData(ext_len, 2)) break;
            payload_len = (static_cast<uint8_t>(ext_len[0]) << 8) | static_cast<uint8_t>(ext_len[1]);
        } else if (payload_len == 127) {
            char ext_len[8];
            if (!readData(ext_len, 8)) break;
            payload_len = 0;
            for (int i = 0; i < 8; ++i) {
                payload_len = (payload_len << 8) | static_cast<uint8_t>(ext_len[i]);
            }
        }
        
        // 读取masking key
        char masking_key[4];
        if (masked) {
            if (!readData(masking_key, 4)) break;
        }
        
        // 读取payload
        std::vector<uint8_t> payload(payload_len);
        if (payload_len > 0 && !readData(reinterpret_cast<char*>(payload.data()), payload_len)) {
            break;
        }
        
        // 解码payload
        if (masked) {
            for (size_t i = 0; i < payload.size(); ++i) {
                payload[i] ^= static_cast<uint8_t>(masking_key[i % 4]);
            }
        }
        
        // 处理消息
        std::string message(payload.begin(), payload.end());
        if (message_handler_) {
            message_handler_(message);
        }
        
        std::cout << "收到消息: " << message << std::endl;
    }
    
    connected_ = false;
}

bool WebSocketClient::sendData(const char* data, size_t length) {
    size_t total_sent = 0;
    while (total_sent < length) {
        ssize_t bytes_sent = send(socket_, data + total_sent, length - total_sent, 0);
        if (bytes_sent <= 0) {
            return false;
        }
        total_sent += bytes_sent;
    }
    return true;
}

bool WebSocketClient::readData(char* buffer, size_t size) {
    size_t total_read = 0;
    while (total_read < size) {
        ssize_t bytes_read = recv(socket_, buffer + total_read, size - total_read, 0);
        if (bytes_read <= 0) {
            return false;
        }
        total_read += bytes_read;
    }
    return true;
}

} // namespace websocket